import os
import numpy as np
import paddle
from paddle.io import Dataset
from paddle.vision.transforms import Compose, Resize, Normalize, Transpose, RandomHorizontalFlip
from PIL import Image  # 使用PIL库替代paddle.vision.io


class Garbages_Dataset(Dataset):
    def __init__(self, data_root, list_file, mode="train", transform=None):
        super(Garbages_Dataset, self).__init__()
        self.data_root = data_root
        self.list_file = list_file
        self.mode = mode
        self.transform = transform

        # 读取标签文件，建立类别名到ID的映射
        self.labels = []
        with open(os.path.join(data_root, "labels.txt"), "r", encoding="utf-8") as f:
            for line in f:
                self.labels.append(line.strip())
        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.labels)}

        # 读取数据列表文件（train_list.txt或test_list.txt）
        self.samples = []
        with open(list_file, "r", encoding="utf-8") as f:
            for line in f:
                path, label = line.strip().split()
                self.samples.append((path, int(label)))

    def __getitem__(self, index):
        img_path, label = self.samples[index]

        # 使用PIL打开图像
        img = Image.open(img_path).convert('RGB')  # 确保图像为RGB格式

        # 应用数据变换
        if self.transform:
            img = self.transform(img)

        return img, label  # 返回处理后的图像和标签

    def __len__(self):
        return len(self.samples)

    def get_labels(self):
        return self.labels  # 返回类别名列表


# 训练集变换（包含数据增强）
train_transform = Compose([
    Resize(size=(224, 224)),  # 调整尺寸
    RandomHorizontalFlip(prob=0.5),  # 随机水平翻转
    Normalize(  # 归一化（ImageNet标准）
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
        data_format='HWC'
    ),
    Transpose()  # 确保格式为CHW（Paddle默认）
])

# 测试集变换（仅预处理，无数据增强）
test_transform = Compose([
    Resize(size=(224, 224)),
    Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
        data_format='HWC'
    ),
    Transpose()
])

# 数据集路径
data_root = "/home/aistudio/data/data89094/Garbages"
train_list = os.path.join(data_root, "train_list.txt")
test_list = os.path.join(data_root, "test_list.txt")

# 创建数据集实例
train_dataset = Garbages_Dataset(
    data_root=data_root,
    list_file=train_list,
    mode="train",
    transform=train_transform
)

test_dataset = Garbages_Dataset(
    data_root=data_root,
    list_file=test_list,
    mode="test",
    transform=test_transform
)

# 创建数据加载器
train_loader = paddle.io.DataLoader(
    train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

test_loader = paddle.io.DataLoader(
    test_dataset,
    batch_size=32,
    shuffle=False,
    num_workers=4
)

# 构建模型（以ResNet50为例）
model = paddle.vision.models.resnet50(pretrained=True)

# 获取原始模型fc层的输入维度（通过.weight.shape[1]）
in_features = 2048  # 替代in_features

# 修改最后一层为自定义类别数
model.fc = paddle.nn.Linear(in_features, len(train_dataset.get_labels()))

# 转换为Paddle Model（用于高层API训练）
model = paddle.Model(model)

# 配置训练参数
model.prepare(
    optimizer=paddle.optimizer.Adam(parameters=model.parameters(), learning_rate=0.001),
    loss=paddle.nn.CrossEntropyLoss(),
    metrics=paddle.metric.Accuracy()
)

# 启动训练
model.fit(
    train_loader,
    test_loader,
    epochs=10,
    batch_size=32,
    verbose=1,
    save_dir="output/garbage_classifier"
)

# 评估模型
eval_result = model.evaluate(test_loader)
print(f"评估指标：{eval_result}")

model.save("Model/Garbage01", training=False)